#include "VideoEncoder.h"
#include "FpsCalculator.h"
#include "Poller/EventPoller.h"
#include "VideoDecoder.h"
#include "spdlog/spdlog.h"
#include <asm-generic/errno-base.h>
#include <chrono>

#include <libavcodec/avcodec.h>
#include <libavcodec/codec.h>
#include <libavcodec/packet.h>
#include <libavformat/avformat.h>
#include <libavformat/avio.h>
#include <libavutil/buffer.h>
#include <libavutil/error.h>
#include <libavutil/frame.h>
#include <libavutil/hwcontext.h>
#include <libavutil/mathematics.h>
#include <libavutil/pixfmt.h>
#include <memory>
#include <mutex>
#include <opencv2/core/cuda.hpp>
#include <opencv2/core/hal/interface.h>
#include <thread>

#include "cuda_util.h"

void VideoEncoder::Init(int fps, bool h264, const std::string &rtsp_url, const int32_t gpu_idx, const int32_t core_idx)
{
    _fps = fps;
    _h264 = h264;
    _rtsp_url = rtsp_url;
    _gpu_index = gpu_idx;
    _core_idx = core_idx;

    _circular_queue.set_capacity(5);

    _encode_fps = std::make_shared<FpsCalculator>();
    _encode_fps->Init(50, "encode", true);

    _npp_stream_ctx.hStream = nullptr;
    _npp_stream_ctx.nCudaDeviceId = _gpu_index;
    _npp_stream_ctx.nMultiProcessorCount = 0;
    _npp_stream_ctx.nMaxThreadsPerMultiProcessor = 0;
    _npp_stream_ctx.nReserved0 = 0;

    cudaStreamCreate(&_npp_stream_ctx.hStream);

    _last_time = std::chrono::high_resolution_clock::now();
}

void VideoEncoder::Start()
{
    _start = true;

    if (_circular_queue.empty())
    {
        toolkit::EventPollerPool::Instance().getPoller()->doDelayTask(2000, [this]() {
            Start();
            return 0;
        });
    }
    else
    {
        auto image = _circular_queue.front();
        _width = image->width;
        _height = image->height;

        _encode_thread = std::make_shared<std::thread>([this]() {
            while (_start)
            {
                if (_init())
                {
                    EncodeThread();
                }
                _free();

                if (_start)
                {
                    std::this_thread::sleep_for(3s);
                }
            }
        });
    }
}

void VideoEncoder::Stop()
{
    _start = false;
    _cv.notify_one();

    if (_encode_thread)
    {
        _encode_thread->join();
        _encode_thread = nullptr;
    }
}

void VideoEncoder::AddImage(std::shared_ptr<ImageInfo> image)
{
    {
        std::scoped_lock<std::mutex> g(_mtx);
        _circular_queue.push_back(image);
    }
    _cv.notify_one();
}

VideoEncoder::~VideoEncoder()
{
    if (_npp_stream_ctx.hStream)
        cudaStreamDestroy(_npp_stream_ctx.hStream);
}

bool VideoEncoder::_init()
{
    avformat_alloc_output_context2(&_fmt_ctx, nullptr, "rtsp", _rtsp_url.c_str());
    if (!_fmt_ctx)
    {
        SPDLOG_ERROR("avformat_alloc_output_context2 error");
        return false;
    }

    AVCodec *find_codec = nullptr;
    if (_h264)
    {
        auto codec = avcodec_find_encoder_by_name("h264_nvenc");
        if (!codec)
        {
            SPDLOG_ERROR("Failed to find h264_nvenc encoder, fallback to x264");
            codec = avcodec_find_encoder(AV_CODEC_ID_H264);
        }
        find_codec = (AVCodec *)codec;
    }
    else
    {
        auto codec = avcodec_find_encoder_by_name("hevc_nvenc");
        if (!codec)
        {
            SPDLOG_ERROR("Failed to find h264_nvenc encoder, fallback to h265");
            codec = avcodec_find_encoder(AV_CODEC_ID_HEVC);
        }
        find_codec = (AVCodec *)codec;
    }

    if (!find_codec)
    {
        SPDLOG_ERROR("can't find encoder");
        return false;
    }

    _codec_ctx = avcodec_alloc_context3(find_codec);
    _codec_ctx->width = _width;
    _codec_ctx->height = _height;
    _codec_ctx->time_base = AVRational { 1, _fps };
    _codec_ctx->framerate = AVRational { _fps, 1 };
    _codec_ctx->gop_size = _fps;
    _codec_ctx->max_b_frames = 0;
    _codec_ctx->pix_fmt = AV_PIX_FMT_CUDA;

    auto hw_device_ctx = SharedResource::GetInstance()->HW_Device_Ctx(_gpu_index);
    _codec_ctx->hw_device_ctx = av_buffer_ref(hw_device_ctx);

    _hw_frames_ref = av_hwframe_ctx_alloc(hw_device_ctx);

    if (_hw_frames_ref == nullptr)
    {
        SPDLOG_ERROR("Failed to alloc hwframe_ctx");
        return false;
    }

    _frames_ctx = (AVHWFramesContext *)_hw_frames_ref->data;
    _frames_ctx->format = AV_PIX_FMT_CUDA;
    _frames_ctx->sw_format = AV_PIX_FMT_NV12;
    _frames_ctx->width = _width;
    _frames_ctx->height = _height;
    _frames_ctx->initial_pool_size = 25;

    if (av_hwframe_ctx_init(_hw_frames_ref))
    {
        SPDLOG_ERROR("Failed to init hwframe_ctx");
        return false;
    }

    _codec_ctx->hw_frames_ctx = av_buffer_ref(_hw_frames_ref);

    // 设置nvenc特定参数
    // av_opt_set(_codec_ctx->priv_data, "preset", "llhp", 0); // low-latency high performance
    // av_opt_set(_codec_ctx->priv_data, "tune", "zerolatency", 0);

    if (_fmt_ctx->oformat->flags & AVFMT_GLOBALHEADER)
        _codec_ctx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;

    if (avcodec_open2(_codec_ctx, find_codec, nullptr) < 0)
    {
        SPDLOG_ERROR("Failed to open codec");
        return false;
    }

    _video_stream = avformat_new_stream(_fmt_ctx, nullptr);
    if (!_video_stream)
    {
        SPDLOG_ERROR("avformat_new_stream error");
        return false;
    }

    _video_stream->time_base = _codec_ctx->time_base;
    avcodec_parameters_from_context(_video_stream->codecpar, _codec_ctx);

    if (!(_fmt_ctx->oformat->flags & AVFMT_NOFILE))
    {
        if (avio_open(&_fmt_ctx->pb, _rtsp_url.c_str(), AVIO_FLAG_WRITE) < 0)
        {
            SPDLOG_ERROR("Failed to open output IO");
            return false;
        }
    }

    if (avformat_write_header(_fmt_ctx, nullptr) < 0)
    {
        SPDLOG_ERROR("Failed to write header");
        return false;
    }

    _pkt = av_packet_alloc();

    return true;
}

void VideoEncoder::_free()
{
    if (_codec_ctx)
    {
        // 刷新编码器缓冲区
        int ret = avcodec_send_frame(_codec_ctx, nullptr);
        if (ret < 0 && ret != AVERROR_EOF)
        {
            SPDLOG_WARN("Error sending flush frame: {}", ret);
        }
        else
        {
            while (ret >= 0)
            {
                ret = avcodec_receive_packet(_codec_ctx, _pkt);
                if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
                {
                    break;
                }
                else if (ret < 0)
                {
                    SPDLOG_WARN("avcodec_receive_packet failed during flush: {}", ret);
                    break;
                }

                _pkt->stream_index = _video_stream->index;
                _pkt->pts = av_rescale_q(_pkt->pts, _codec_ctx->time_base, _video_stream->time_base);
                _pkt->dts = _pkt->pts;

                ret = av_interleaved_write_frame(_fmt_ctx, _pkt);
                if (ret < 0)
                {
                    SPDLOG_WARN("av_interleaved_write_frame failed during flush: {}", ret);
                }
                av_packet_unref(_pkt);
            }
        }

        if (av_write_trailer(_fmt_ctx) < 0)
        {
            SPDLOG_WARN("Error writing trailer");
        }
    }

    if (_codec_ctx)
    {
        avcodec_free_context(&_codec_ctx);
    }

    if (_pkt)
    {
        av_packet_free(&_pkt);
    }

    if (_fmt_ctx)
    {
        if (_fmt_ctx && !(_fmt_ctx->oformat->flags & AVFMT_NOFILE))
        {
            avio_close(_fmt_ctx->pb);
        }
        avformat_free_context(_fmt_ctx);
    }

    if (_hw_frames_ref)
    {
        av_buffer_unref(&_hw_frames_ref);
    }
}

void VideoEncoder::EncodeThread()
{
    while (_start)
    {
        cv::cuda::GpuMat bgr;
        bool got_frame = false;

        {
            std::unique_lock<std::mutex> g(_mtx);
            if (_cv.wait_for(g, std::chrono::milliseconds(40), [this] { return !_circular_queue.empty(); }))
            {
                auto image = _circular_queue.front();
                _circular_queue.pop_front();
                bgr = image->image;
                got_frame = true;
                _last_image = image;
            }
        }

        if (!got_frame)
        {
            // TODO: 超时没有新数据
            if (_last_image)
            {
                bgr = _last_image->image;
            }
            else
            {
                continue;
            }
        }

        _encode_fps->Add();

        {
            NppiSize roi = { _width, _height };

            cv::cuda::GpuMat y_plane(_height, _width, CV_8UC1);
            cv::cuda::GpuMat uv_plane(_height / 2, _width / 2, CV_8UC2);

            Npp8u *src = bgr.ptr<Npp8u>();
            int src_step = static_cast<int>(bgr.step);

            BGR_to_NV12(
                bgr.ptr<Npp8u>(), y_plane.ptr<Npp8u>(), uv_plane.ptr<Npp8u>(), static_cast<int>(bgr.step), static_cast<int>(y_plane.step),
                static_cast<int>(uv_plane.step), _width, _height, _npp_stream_ctx.hStream);

            AVFrame *frame = av_frame_alloc();
            frame->format = AV_PIX_FMT_NV12;
            frame->width = _width;
            frame->height = _height;
            frame->hw_frames_ctx = av_buffer_ref(_codec_ctx->hw_frames_ctx);
            frame->pts = av_rescale_q(_frame_idx, { 1, _fps }, _codec_ctx->time_base);

            if (av_hwframe_get_buffer(_hw_frames_ref, frame, 0) < 0)
            {
                SPDLOG_ERROR("Failed to get gpu frame buffer");
                av_frame_unref(frame);
                av_frame_free(&frame);
                continue;
            }

            cudaMemcpy2DAsync(
                frame->data[0], frame->linesize[0], y_plane.ptr(), y_plane.step, _width, _height, cudaMemcpyDeviceToDevice, _npp_stream_ctx.hStream);
            cudaMemcpy2DAsync(
                // frame->data[1], frame->linesize[1], uv_plane.ptr(), uv_plane.step, _width, _height / 2, cudaMemcpyDeviceToDevice, _npp_stream_ctx.hStream);
                frame->data[1], frame->linesize[1], uv_plane.ptr(), _width, _width, _height / 2, cudaMemcpyDeviceToDevice, _npp_stream_ctx.hStream);

            // frame->data[0] = y_plane.ptr();
            // frame->linesize[0] = y_plane.step;

            // frame->data[1] = uv_plane.ptr();
            // frame->linesize[1] = uv_plane.step;

            cudaStreamSynchronize(_npp_stream_ctx.hStream);
            int ret = avcodec_send_frame(_codec_ctx, frame);

            if (ret < 0)
            {
                if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
                {
                    continue;
                }
                else
                {
                    break;
                }
            }

            _frame_idx++;

            while (ret >= 0)
            {
                ret = avcodec_receive_packet(_codec_ctx, _pkt);
                if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
                {
                    break;
                }
                else if (ret < 0)
                {
                    SPDLOG_ERROR("avcodec_receive_packet failed: {}", ret);
                    break;
                }
                _pkt->stream_index = _video_stream->index;
                _pkt->pts = av_rescale_q(_pkt->pts, _codec_ctx->time_base, _video_stream->time_base);
                _pkt->dts = _pkt->pts;

                ret = av_interleaved_write_frame(_fmt_ctx, _pkt);

                if (ret < 0)
                {
                    SPDLOG_ERROR("av_interleaved_write_frame failed: {}", ret);
                    av_packet_unref(_pkt);
                    break;
                }

                av_packet_unref(_pkt);
            }

            av_frame_free(&frame);
        }

        static auto next_frame_time = std::chrono::high_resolution_clock::now();
        std::this_thread::sleep_until(next_frame_time);
        next_frame_time += std::chrono::milliseconds(1000 / _fps);

        // auto current_time = std::chrono::high_resolution_clock::now();
        // auto diff = std::chrono::duration_cast<std::chrono::milliseconds>(current_time - _last_time).count();
        // if (diff < (1000 / _fps))
        // {
        //     auto interval = (1000 / _fps) - diff;
        //     SPDLOG_INFO("Diff : {}   ,   Speep for: {} ms", diff, interval);
        //     std::this_thread::sleep_for(std::chrono::milliseconds(interval ));
        // }
        // else
        // {
        //     // std::this_thread::sleep_for(std::chrono::milliseconds(200));
        //     SPDLOG_INFO("Diff : {}", diff);
        // }
        // _last_time = current_time;
    }
}